import websocket  # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid
import json
import requests
import os
import sys
from datetime import datetime

server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())

def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    headers = {'Content-Type': 'application/json'}
    response = requests.post(f"http://{server_address}/prompt", json=p, headers=headers)
    response.raise_for_status()  # 如果返回错误状态码，则抛出异常
    return response.json()

def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    response = requests.get(f"http://{server_address}/view", params=data)
    response.raise_for_status()
    return response.content

def get_history(prompt_id):
    response = requests.get(f"http://{server_address}/history/{prompt_id}")
    response.raise_for_status()
    return response.json()

def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break  # Execution is done
        else:
            continue  # previews are binary data

    history = get_history(prompt_id)[prompt_id]
    for o in history['outputs']:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
            output_images[node_id] = images_output

    return output_images

if len(sys.argv) != 6:
    print("Usage: python api2.py <clothes_image_path> <inspire_ink_image_path> <output_directory> <json_file_path> <mask_file_path>")
    sys.exit(1)

clothes_image_path = sys.argv[1]
inspire_ink_image_path = sys.argv[2]
output_directory = sys.argv[3]
json_file_path = sys.argv[4]
mask_file_path = sys.argv[5]

print(f"Clothes Image Path: {clothes_image_path}")
print(f"Inspire Ink Image Path: {inspire_ink_image_path}")
print(f"Output Directory: {output_directory}")
print(f"JSON File Path: {json_file_path}")
print(f"Mask File Path: {mask_file_path}")

with open(json_file_path, "r", encoding="utf-8") as f:
    workflow_jsondata = f.read()

prompt = json.loads(workflow_jsondata)
prompt["117"]["inputs"]["image"] = clothes_image_path
prompt["306"]["inputs"]["image"] = inspire_ink_image_path
prompt["118"]["inputs"]["image"] = mask_file_path  # 使用新的 mask 文件路径

ws = websocket.WebSocket()
ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
images = get_images(ws, prompt)

current_time = datetime.now().strftime("%Y%m%d%H%M%S")

# 展示输出图像
for node_id in images:
    if node_id in ["128", "135"]:
        for image_data in images[node_id]:
            from PIL import Image
            import io
            image = Image.open(io.BytesIO(image_data))
            image.show()
            if node_id == "128":
                image_path = os.path.join(output_directory, f"clotheswithback_{current_time}.png")
            elif node_id == "135":
                image_path = os.path.join(output_directory, f"clotheswithoutback_{current_time}.png")
            print(f"Saving image at: {image_path}")
            image.save(image_path)